import asyncio
import websockets
import json
import base64
from pathlib import Path

async def handle_client(websocket):
    buffer = bytearray()
    banner_processed = False
    frame_body_length = 0
    frame_count = 0

    # Banner解析状态
    current_banner = {
        'version': 0,
        'length': 0,
        'pid': 0,
        'realWidth': 0,
        'realHeight': 0,
        'virtualWidth': 0,
        'virtualHeight': 0,
        'orientation': 0,
        'quirks': 0
    }

    async for message in websocket:
        buffer.extend(message)

        while True:
            # 解析Banner
            if not banner_processed:
                if len(buffer) < 2:
                    break  # 等待更多数据

                # 读取banner版本和长度
                version = buffer[0]
                banner_length = buffer[1]
                current_banner['version'] = version
                current_banner['length'] = banner_length

                if len(buffer) < banner_length:
                    break  # 数据不足

                # 提取banner字节并解析
                banner_bytes = bytes(buffer[:banner_length])
                buffer = buffer[banner_length:]

                # 解析各字段（小端序处理）
                current_banner['pid'] = int.from_bytes(banner_bytes[2:6], 'little')
                current_banner['realWidth'] = int.from_bytes(banner_bytes[6:10], 'little')
                current_banner['realHeight'] = int.from_bytes(banner_bytes[10:14], 'little')
                current_banner['virtualWidth'] = int.from_bytes(banner_bytes[14:18], 'little')
                current_banner['virtualHeight'] = int.from_bytes(banner_bytes[18:22], 'little')
                current_banner['orientation'] = banner_bytes[22] * 90
                current_banner['quirks'] = banner_bytes[23]

                # 保存banner信息
                banner_info = {
                    "raw_data": base64.b64encode(banner_bytes).decode(),
                    "parsed": current_banner
                }
                output_dir = Path('output')
                output_dir.mkdir(parents=True, exist_ok=True)
                with open('output/banner.json', 'w') as f:
                    json.dump(banner_info, f, indent=2)
                
                banner_processed = True
                print("Banner processed")

            # 解析帧数据
            else:
                if frame_body_length == 0:
                    if len(buffer) < 4:
                        break  # 需要更多数据
                    frame_body_length = int.from_bytes(buffer[:4], 'little')
                    buffer = buffer[4:]

                if len(buffer) < frame_body_length:
                    break  # 数据不足

                frame_data = bytes(buffer[:frame_body_length])
                buffer = buffer[frame_body_length:]

                # 验证JPEG头
                if frame_data[0] != 0xFF or frame_data[1] != 0xD8:
                    print("Invalid JPEG header")
                    continue

                # 保存图片
                frame_count += 1
                Path(f'output/frame_{frame_count}.jpg').write_bytes(frame_data)
                print(f"Saved frame {frame_count}")

                frame_body_length = 0  # 重置状态

async def main():
    async with websockets.serve(handle_client, "localhost", 9002):
        await asyncio.Future() 

if __name__ == "__main__":
    asyncio.run(main())